from langchain_core.chat_models.base import BaseChatModel
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from typing import Any, List, Optional, Dict
from deepseek import DeepSeekAPI
import logging

logger = logging.getLogger(__name__)

class DeepSeekChatModel(BaseChatModel):
    """自定义的 DeepSeek Chat Model"""
    
    def __init__(self):
        super().__init__()
        self.deepseek_api = DeepSeekAPI()
    
    @property
    def _llm_type(self) -> str:
        """返回 LLM 类型"""
        return "deepseek_chat"

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """生成回复"""
        try:
            # 将消息转换为文本
            prompt = self._convert_messages_to_prompt(messages)
            logger.info(f"Sending prompt to DeepSeek API: {prompt[:200]}...")
            
            # 调用 DeepSeek API
            response = self.deepseek_api.chat_completion(prompt)
            logger.info(f"Received response from DeepSeek API: {response[:200]}...")
            
            return {
                "generations": [{
                    "text": response,
                    "message": AIMessage(content=response),
                }]
            }
        except Exception as e:
            logger.error(f"Error in DeepSeek API call: {str(e)}")
            raise
    
    def _convert_messages_to_prompt(self, messages: List[BaseMessage]) -> str:
        """将消息列表转换为提示词"""
        prompt_parts = []
        for message in messages:
            if isinstance(message, HumanMessage):
                prompt_parts.append(f"Human: {message.content}")
            elif isinstance(message, AIMessage):
                prompt_parts.append(f"Assistant: {message.content}")
            else:
                prompt_parts.append(message.content)
        return "\n".join(prompt_parts)

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """返回模型参数"""
        return {
            "model_type": "deepseek_chat",
        } 